package com.lu.springbootinit.manager;

import com.lu.springbootinit.common.ErrorCode;
import com.lu.springbootinit.exception.BusinessException;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChoice;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Component
public class AiManager {
    @Resource
    private ArkService service;

    private final String DEFAULT_MODEL = "deepseek-r1-250120";

    /**
     *  调用 AI 获取 AI 响应字符串
     * @param
     */
    public String doChat( String userPromt) {
        return doChat("", userPromt);
    }
    /**
     *  调用 AI 获取 AI 响应字符串
     * @param
     */
    public String doChat(String systemPromt, String userPromt) {
        return doChat(systemPromt, userPromt, DEFAULT_MODEL);
    }
    /**
     *  调用 AI 获取 AI 响应字符串
     * @param
     */
    public String doChat(String systemPromt, String userPromt, String modelName) {
        final List<ChatMessage> streamMessages = new ArrayList<>();
        // 构造消息
        final ChatMessage systemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content(systemPromt).build();
        final ChatMessage streamUserMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userPromt).build();

        streamMessages.add(systemMessage);
        streamMessages.add(streamUserMessage);
        return doChat(streamMessages, modelName);
    }

    /**
     * 调用 AI 接口，获取响应字符串（允许传入自定义的消息列表，使用默认模型）
     *
     * @param messages
     * @return
     */
    public String doChat(List<ChatMessage> messages) {
        return doChat(messages, DEFAULT_MODEL);
    }

    /**
     * 调用 AI 接口，获取响应字符串（允许传入自定义的消息列表）
     *
     * @param messages
     * @param model
     * @return
     */
    public String doChat(List<ChatMessage> messages, String model) {
        // 构造请求
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(messages)
                .build();
        // 调用接口发送请求
        List<ChatCompletionChoice> choices = service.createChatCompletion(streamChatCompletionRequest).getChoices();
        if(choices != null) {
            return (String) choices.get(0).getMessage().getContent();
        }else {
            throw new BusinessException(ErrorCode.OPERATION_ERROR, "AI 调用失败,没有返回结果");
        }
    }
    /**
     *  调用 AI 获取 AI 响应字符串（流式输出结果）
     *  （无使用）
     * @param
     */
    public void doChatFlow(String systemPromt, String userPromt, String modelName) {
        final List<ChatMessage> streamMessages = new ArrayList<>();
        // 构造消息
        final ChatMessage systemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content(systemPromt).build();
        final ChatMessage streamUserMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userPromt).build();

        streamMessages.add(systemMessage);
        streamMessages.add(streamUserMessage);

        // 构造请求
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelName)
                .messages(streamMessages)
                .build();
        // 调用接口发送请求 (流式输出结果)
        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(
                        delta -> {
                            if (!delta.getChoices().isEmpty()) {
                                if (StringUtils.isNotEmpty(delta.getChoices().get(0).getMessage().getReasoningContent())) {
                                    System.out.print(delta.getChoices().get(0).getMessage().getReasoningContent());
                                } else {
                                    System.out.print(delta.getChoices().get(0).getMessage().getContent());
                                }
                            }
                        }
                );
        service.shutdownExecutor();
    }


}
